
import os,sys

osa = os.path.abspath
osd = os.path.dirname
cur_dir = osd(osa(__file__))
par_dir = osd(cur_dir)
sys.path.insert(0,par_dir)



import os
import json
from tqdm import tqdm
import shutil
from PIL import Image

from utils.util_for_os import osj,ose

# /mnt/nas/shengjie/datasets/dataset_depth_control_depth/
# control: ori depth
# file: restored depth
'''
1, read names.txt
2, 前 1450 for train
   后 ~end for val
3, 读取 一张图                                  for image
   替换这张图 最后 1/3 为(2/3w,0,w,h) 为 空白    for control
   分别 save (train | val)/(image | control_image)

   prompt = [depredux] merge (shape of left side and style of middle side) to right side
'''

depredux_dir = '/mnt/nas/shengjie/datasets/dataset_depth_control_1364'
depredux_names_path = osj(depredux_dir, 'names.txt')
# 读取 names.txt，筛选所有的jpg图片
with open(depredux_names_path, "r") as f:
    all_names = [line.strip() for line in f if line.strip()]

# 只保留.jpg结尾的文件名
jpg_names = [name for name in all_names if name.lower().endswith('.jpg')]

# 新数据集路径
base_dir = "/mnt/nas/shengjie/datasets/KontextRefControl_depredux"
os.makedirs(base_dir, exist_ok=True)

# 创建子目录
train_image_dir = os.path.join(base_dir, "train/image")
train_control_dir = os.path.join(base_dir, "train/control_image")
val_image_dir = os.path.join(base_dir, "val/image")
val_control_dir = os.path.join(base_dir, "val/control_image")

os.makedirs(train_image_dir, exist_ok=True)
os.makedirs(train_control_dir, exist_ok=True)
os.makedirs(val_image_dir, exist_ok=True)
os.makedirs(val_control_dir, exist_ok=True)


# 固定 prompt
prompt = "[depredux] merge (shape of left side and style of middle side) to right side"

# 训练/验证划分
train_num = 1355
train_names = jpg_names[:train_num]
val_names = jpg_names[train_num:]

def process_depredux_split(names, image_dir, control_dir, meta_path):
    meta_data = []
    for name in tqdm(names, desc=f"Processing {meta_path}"):
        img_path = osj(depredux_dir, name)
        img = Image.open(img_path)
        w, h = img.size
        # 将图像的最后1/3区域(2/3w, 0, w, h)替换为空白
        blank_box = (int(w * 2 / 3), 0, w, h)
        img_blank = img.copy()
        # 创建与原图相同模式的白色区域
        if img_blank.mode == "RGBA":
            white = (255, 255, 255, 255)
        else:
            white = (255, 255, 255)
        from PIL import ImageDraw
        draw = ImageDraw.Draw(img_blank)
        draw.rectangle(blank_box, fill=white)
        control_img = img_blank

        # 保存文件名
        stem = os.path.splitext(name)[0]
        ext = os.path.splitext(name)[1]
        image_name = f"{stem}_image{ext}"
        control_name = f"{stem}_control{ext}"

        # 保存图片
        control_img.save(osj(control_dir, control_name)) 
        img.save(osj(image_dir, image_name))

        # 记录元数据
        meta_entry = {
            "file_name": f"image/{image_name}",
            "control_image": f"control_image/{control_name}",
            "prompt": prompt
        }
        meta_data.append(meta_entry)
    # 写入jsonl
    with open(meta_path, "w") as f:
        for item in meta_data:
            f.write(json.dumps(item) + "\n")

# 处理训练集
process_depredux_split(
    train_names,
    train_image_dir,
    train_control_dir,
    osj(base_dir, "train/metadata.jsonl")
)

# 处理验证集
process_depredux_split(
    val_names,
    val_image_dir,
    val_control_dir,
    osj(base_dir, "val/metadata.jsonl")
)

# 生成数据集描述文件
dataset_info = {
    "train_size": len(train_names),
    "val_size": len(val_names),
    "resolution": "3072*1024",
    "default_prompt": prompt
}
with open(osj(base_dir, "dataset_info.json"), "w") as f:
    json.dump(dataset_info, f, indent=2)

print("✅ JSONL 数据集生成完成！")
